Codecov Report

Attention: Patch coverage is 92.00000% with 6 lines in your changes missing coverage. Please review.

Project coverage is 78.50%. Comparing base (d3f22b5) to head (ce7c656).
Report is 15 commits behind head on main.

@@             Coverage Diff             @@
##             main    #1269       +/-   ##
===========================================
- Coverage   89.40%   78.50%   -10.91%     
===========================================
  Files         118      119        +1     
  Lines        8715     8811       +96     
===========================================
- Hits         7792     6917      -875     
- Misses        923     1894      +971     

Hey @janfb,
would very much appreciate your input at this stage:

Currently the PR adds the CategoricalMADE and builder build_autoregressive_categoricalestimator + some minor modifications to build_mnle and MixedDensityEstimator. This enables multiple discrete variables with different numbers of classes via trainer = MNLE(density_estimator=lambda x,y: build_mnle(y,x,categorical_model="made")) Note that for some reason x and y have to be flipped for mnle.

As far as I can tell all functionalities of CategoricalMADE work for both 1D and ND inputs and running the Example_01_DecisionMakingModel.ipynb with the CatMADE matches the ground truth

The question now is: How should I verify this works? / Which tests should I add/modify? Do you have an idea for a good toy example with several discrete variables that I could use?

I have cooked up a toy simulator, for which I am getting good posteriors using SNPE, but for some reason MNLE raises a RuntimeError: probability tensor contains either 'inf', 'nan' or element < 0 Even for the unmodified MNLE. Any ideas why this could be?

This is the simulator

def toy_simulator(theta: torch.Tensor, centers: list[torch.Tensor]) -> torch.Tensor:
    batch_size, n_dimensions = theta.shape
    assert len(centers) == n_dimensions, "Number of center sets must match theta dimensions"
    
    # Calculate discrete classes by assiging to the closest center
    x_disc = torch.stack([
        torch.argmin(torch.abs(centers[i].unsqueeze(1) - theta[:, i].unsqueeze(0)), dim=0)
        for i in range(n_dimensions)
    ], dim=1)

    closest_centers = torch.stack([centers[i][x_disc[:, i]] for i in range(n_dimensions)], dim=1)
    # Add Gaussian noise to assigned class centers
    std = 0.4
    x_cont = closest_centers + std * torch.randn_like(closest_centers)
       
    return torch.cat([x_cont, x_disc], dim=1)

The setup:

torch.random.manual_seed(0)
centers = [
    torch.tensor([-0.5, 0.5]),
    # torch.tensor([-1.0, 0.0, 1.0]),
]

prior = BoxUniform(low=torch.tensor([-2.0]*len(centers)), high=torch.tensor([2.0]*len(centers)))
theta = prior.sample((20000,))
x = toy_simulator(theta, centers)

theta_o = prior.sample((1,))
x_o = toy_simulator(theta_o, centers)

NPE:

trainer = SNPE()
estimator = trainer.append_simulations(theta=theta, x=x).train(training_batch_size=1000)

snpe_posterior = trainer.build_posterior(prior=prior)
posterior_samples = snpe_posterior.sample((2000,), x=x_o)
pairplot(posterior_samples, limits=[[-2, 2], [-2, 2]], figsize=(5, 5), points=theta_o)

and the equivalent MNLE:

trainer = MNLE()
estimator = trainer.append_simulations(theta=theta, x=x).train(training_batch_size=1000)

mnle_posterior = trainer.build_posterior(prior=prior)
mnle_samples = mnle_posterior.sample((10000,), x=x_o)
pairplot(mnle_samples, limits=[[-2, 2], [-2, 2]], figsize=(5, 5), points=theta_o)

Hoping this makes sense. Lemme know if you need clarifications anywhere. Thanks for your feedback.

Hey @janfb,
you might have missed this, but I would be happy about feedback :)

Cool, thanks for all the feedback! A quick call would be great, also to discuss suitable tests for this. Will reach out via email and tackle the straight forward things in the meantime.

After discussion with @janfb I will:

1. adapt the simulator of Example_01_DecisionMakingModel.ipynb to multiple discrete variables.
2. Get this to run for 1D and ND
3. fix remaining comments/issues
4. possibly refactor (fold 1D into ND code).

@janfb could you still check tho what is up with the simulator above? Do you have a hunch why the SNPE and MNLE posteriors different?

EDIT:

1. wip
2. ✅
3. ✅
4. ✅
5. add new tests / update old ones with mutli dim example.

I did a bit more work on this PR, current tests should be passing and I have swapped out all the legacy CategoricalNet code for the CategoricalMADE. See changes and comments above (please close if no longer relevant).
A few things remain:

• Add a bit more docs and comments
• add test cases
• adapt the tutorial to 2d? (Can just add another beta distribution to the prior)
• make sure it runs for ND.

This last thing has been haunting me in my sleep, as I cannot figure out what is wrong. Maybe you have an idea of what could be causing this.
For 1D it works, but for ND it always gets the first discrete dim wrong, i.e. yields the prior (all other dims are correct, see image). I am not sure if the conditioning for the first dimension is broken somehow, but I am not able to pin down where this would be happening in my code.

Help would be much appreciated. @janfb

Thanks for all the input <3, looking into the remaining ones over the coming days hopefully

Turns out, since the posterior is an MCMCPosterior, only log_prob is used for sampling and not sample. This bug still eludes me.

I have spent some time today and been able to rule a lot of things out (i.e. posterior, sampling, MNLE related things...), which is great, but nonetheless I am still stuck. I have been able to reduce it to the following example of just training a CategoricalMADE, which means that whatever is causing the weird behaviour can be found here. I am sure I am completely missing something probably obvious. I hope the code below makes it easier to identify.
@dgedon , @janfb .

#... snle tutorial (incl in this PR)

from sbi.neural_nets.estimators.categorical_net import CategoricalMADE

# Define independent prior.
prior = MultipleIndependent(
    [
        Gamma(torch.tensor([1.0]), torch.tensor([0.5])),
        Beta(torch.tensor([2.0]), torch.tensor([2.0])),
        Beta(torch.tensor([2.0]), torch.tensor([2.0])),
        # Beta(torch.tensor([2.0]), torch.tensor([2.0])),
    ],
    validate_args=False,
)

torch.manual_seed(42)
theta_o = prior.sample((1,))

# Training data
num_simulations = 10000
batch_size = 1000
num_epochs = 100
theta = prior.sample((num_simulations,))
x = mixed_simulator(theta)

# only pred disc dimensions
x = x[:, 1:]

made = CategoricalMADE(
    num_categories=torch.ones(x.shape[1], dtype=torch.int32)*2,
    hidden_features=20,
    context_features=theta.shape[1],
)

# quick and dirty training loop
in_batches = lambda x: x.reshape(num_simulations // batch_size, batch_size, -1)
optimizer = torch.optim.Adam(made.parameters(), lr=5e-4)
for i in range(num_epochs):
    print(f"\repoch {i+1} / {num_epochs}", end="")
    for theta_batch, x_batch in zip(in_batches(theta), in_batches(x)):
        optimizer.zero_grad()
        loss = -made.log_prob(x_batch, theta_batch).mean()
        loss.backward()
        optimizer.step()

p_true_disc = theta_o[0, 1:] # theta specifies the true probs
num_disc = x.shape[1]

# compute marginal likelihoods p(x)
choices = torch.arange(2**num_disc).unsqueeze(-1).bitwise_and(2**torch.arange(num_disc)).ne(0).unsqueeze(1)
p_est_disc = torch.zeros(num_disc)
for i in range(num_disc):
    ways_of_choosing_i = choices[torch.any(choices[:, :, i], dim=-1)].float()
    log_prob = made.log_prob(ways_of_choosing_i, theta_o)
    p_est_disc[i] = torch.exp(log_prob).sum().detach()

print("\n")
print(f"true: {p_true_disc}")
print(f"est: {p_est_disc}") # <-- dim=0 incorrect for dim_disc > 1